#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient, ASCENDING, DESCENDING
import sys
import HTMLParser
import logging

default_encoding = 'utf-8'
if sys.getdefaultencoding() != default_encoding:
    reload(sys)
    sys.setdefaultencoding(default_encoding)


# 连接MongoDB
def InitMongoDB():
    MONGODB_HOST = 'localhost'
    MONGODB_PORT = 27017
    MONGODB_DB_NAME = 'journal'
    MONGODB_COLLECTION_NAME = 'journal'
    connection = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = connection[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return connection, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# 去空格和标点符号，全小写
def GetStringFeature(line):
    import re
    line = line.lower()
    result = re.sub(r'[^\w]+', '', line)
    return result


# 首先将最大的库Pubmed导入MongoDB
def InsertPubmed(MongoDB_Collection, PubMed_File):
    pubmed_count = 0
    logging.info("Start analyzing PubMed file.")
    infile = open(PubMed_File, 'r')
    outfile = open(PubMed_File + '.tsv', 'w')
    logging.info("PubMed input file: " + PubMed_File)
    logging.info("PubMed output file: " + PubMed_File + ".tsv")
    outfile.write("JrId\tJournalTitle\tMedAbbr\tISSN(Print)\tISSN(Online)\tIsoAbbr\tNlmId\n")
    outfile.flush()
    jrid = ''
    journaltitle = ''
    medabbr = ''
    issnprint = ''
    issnonline = ''
    isoabbr = ''
    nlmid = ''
    insert_record = {}
    infile.readline()
    for line in infile:
        line = line.rstrip("\n")
        if line.startswith("---"):
            outfile.write(
                jrid + "\t" + journaltitle + "\t" + medabbr + "\t" + issnprint + "\t" + issnonline + "\t" + isoabbr + "\t" + nlmid + "\n")
            outfile.flush()
            MongoDB_Collection.insert_one(insert_record)
            pubmed_count += 1
            jrid = ''
            journaltitle = ''
            medabbr = ''
            issnprint = ''
            issnonline = ''
            isoabbr = ''
            nlmid = ''
            insert_record.clear()
        else:
            splitline = line.split(": ")
            if splitline[0] == "JrId":
                jrid = splitline[1]
                if splitline[1] != "":
                    insert_record["JrId"] = int(splitline[1])
            elif splitline[0] == "JournalTitle":
                journaltitle = splitline[1]
                if splitline[1] != "":
                    insert_record["PubMedTitle"] = splitline[1]
                    insert_record["FullTitle"] = [splitline[1]]
            elif splitline[0] == "MedAbbr":
                medabbr = splitline[1]
                if splitline[1] != "":
                    insert_record["PubMedMedAbbr"] = splitline[1]
                    insert_record["MedAbbr"] = [splitline[1]]
            elif splitline[0] == "ISSN (Print)":
                issnprint = splitline[1]
                if splitline[1] != "":
                    insert_record["pISSN"] = splitline[1]
                    insert_record["ISSNSet"] = [splitline[1]]
            elif splitline[0] == "ISSN (Online)":
                issnonline = splitline[1]
                if splitline[1] != "":
                    insert_record["eISSN"] = splitline[1]
                    if splitline[1] != issnprint:
                        if issnprint == '':
                            insert_record["ISSNSet"] = [splitline[1]]
                        else:
                            insert_record["ISSNSet"].append(splitline[1])
            elif splitline[0] == "IsoAbbr":
                isoabbr = splitline[1]
                if splitline[1] != "":
                    insert_record["PubMedISOAbbr"] = splitline[1]
                    insert_record["ISOAbbr"] = [splitline[1]]
            elif splitline[0] == "NlmId":
                nlmid = splitline[1]
                if splitline[1] != "":
                    insert_record["NlmId"] = splitline[1]
            insert_record["PubMed_Record"] = True
    if jrid != '':
        outfile.write(
            jrid + "\t" + journaltitle + "\t" + medabbr + "\t" + issnprint + "\t" + issnonline + "\t" + isoabbr + "\t" + nlmid + "\n")
        outfile.flush()
        MongoDB_Collection.insert_one(insert_record)
        pubmed_count += 1
    infile.close()
    outfile.close()
    MongoDB_Collection.create_index([("NlmId", ASCENDING)])
    MongoDB_Collection.create_index([("pISSN", ASCENDING)])
    MongoDB_Collection.create_index([("eISSN", ASCENDING)])
    MongoDB_Collection.create_index([("ISSNSet", ASCENDING)])
    logging.info(str(pubmed_count) + " records has been handled.")
    logging.info("Finish analyzing PubMed file.")


# 将最相似的PMC库数据导入数据库，使用NlmId和ISSN作为比对方法
def UpsertPMC(MongoDB_Collection, pmc_filepath):
    import csv

    html_parser = HTMLParser.HTMLParser()

    nlm_count = 0
    issn_count = 0
    upsert_count = 0
    logging.info("Start analyzing PMC file.")
    infile = file(pmc_filepath, "rb")
    csvfile = csv.reader(infile)
    csvfile.next()
    for line in csvfile:
        title = line[0].strip()
        title = html_parser.unescape(title)
        nlmta = line[1].strip()
        nlmta = html_parser.unescape(nlmta)
        pISSN = line[2].strip()
        eISSN = line[3].strip()
        publisher = line[4].strip()
        publisher = html_parser.unescape(publisher)
        nlmid = line[5].strip()
        url = line[12].strip()
        result = MongoDB_Collection.find_one({"NlmId": nlmid})
        # NlmId能对上的，直接更新记录
        if result != None:
            nlm_count += 1
            set_ht = {}
            addtoset_ht = {}
            addtoset_issn = []
            if title != "":
                set_ht["PMCTitle"] = title
                addtoset_ht["FullTitle"] = title
            if nlmta != "":
                set_ht["PMCMedAbbr"] = nlmta
                addtoset_ht["MedAbbr"] = nlmta
            if pISSN != "":
                set_ht["pISSN"] = pISSN
                addtoset_issn.append(pISSN)
            if eISSN != "":
                set_ht["eISSN"] = eISSN
                addtoset_issn.append(eISSN)
            if publisher != "":
                set_ht["PMCPublisher"] = publisher
                addtoset_ht["Publisher"] = publisher
            if url != "":
                set_ht["PMC_URL"] = url
                addtoset_ht["URL"] = url
            if len(addtoset_issn) > 0:
                addtoset_ht["ISSNSet"] = {"$each": addtoset_issn}
            set_ht["PMC_Record"] = True
            MongoDB_Collection.update({"NlmId": nlmid}, {"$set": set_ht, "$addToSet": addtoset_ht}, upsert=False)
        # NlmId对不上的，对ISSN
        else:
            issnset = []
            if pISSN != "":
                issnset.append(pISSN)
            if eISSN != "":
                issnset.append(eISSN)
            result2 = MongoDB_Collection.find_one({"ISSNSet": {"$in": issnset}})
            # ISSN能对上的，根据_id更新
            if result2 != None:
                issn_count += 1
                set_ht = {}
                addtoset_ht = {}
                if title != "":
                    set_ht["PMCTitle"] = title
                    addtoset_ht["FullTitle"] = title
                if nlmta != "":
                    set_ht["PMCMedAbbr"] = nlmta
                    addtoset_ht["MedAbbr"] = nlmta
                if publisher != "":
                    set_ht["PMCPublisher"] = publisher
                    addtoset_ht["Publisher"] = publisher
                if url != "":
                    set_ht["PMC_URL"] = url
                    addtoset_ht["URL"] = url
                set_ht["PMC_Record"] = True
                MongoDB_Collection.update({"_id": result2["_id"]}, {"$set": set_ht, "$addToSet": addtoset_ht},
                                          upsert=False)
            # ISSN对不上的，直接插入数据
            else:
                upsert_count += 1
                set_ht = {}
                addtoset_ht = {}
                addtoset_issn = []
                if title != "":
                    set_ht["PMCTitle"] = title
                    addtoset_ht["FullTitle"] = title
                if nlmta != "":
                    set_ht["PMCMedAbbr"] = nlmta
                    addtoset_ht["MedAbbr"] = nlmta
                if pISSN != "":
                    set_ht["pISSN"] = pISSN
                if eISSN != "":
                    set_ht["eISSN"] = eISSN
                if publisher != "":
                    set_ht["PMCPublisher"] = publisher
                    addtoset_ht["Publisher"] = publisher
                if url != "":
                    set_ht["PMC_URL"] = url
                    addtoset_ht["URL"] = url
                if nlmid != "":
                    set_ht["NlmId"] = nlmid
                if len(addtoset_issn) > 0:
                    addtoset_ht["ISSNSet"] = {"$each": addtoset_issn}
                set_ht["PMC_Record"] = True
                MongoDB_Collection.update({"NlmId": nlmid}, {"$set": set_ht, "$addToSet": addtoset_ht}, upsert=True)

    infile.close()
    logging.info("Update " + str(nlm_count) + " records with NlmId.")
    logging.info("Update " + str(issn_count) + " records with ISSN.")
    logging.info("Insert " + str(upsert_count) + " new records.")
    logging.info("Finish analyzing PMC file.")


# 将最相似的JCR库数据导入数据库
def UpsertJCR(MongoDB_Collection, jcr_filepath):
    issn_count = 0
    isoabbr_count = 0
    title_count = 0
    upsert_count = 0
    logging.info("Start analyzing JCR file.")
    infile = open(jcr_filepath, "r")
    infile.readline()
    title_ht = {}
    isoabbr_ht = {}
    iter_db = MongoDB_Collection.find()
    for item in iter_db:
        title_id = item["_id"]
        if item.has_key("PubMedTitle"):
            title_ht[GetStringFeature(item["PubMedTitle"])] = title_id
        if item.has_key("PMCTitle"):
            title_ht[GetStringFeature(item["PMCTitle"])] = title_id
        if item.has_key("PubMedISOAbbr"):
            isoabbr_ht[GetStringFeature(item["PubMedISOAbbr"])] = title_id
    for line in infile:
        line = line.rstrip()
        if line != "":
            splitline = line.split("\t")
            FullTitle = splitline[0].strip()
            ISOAbbr = splitline[1].strip()
            JCRAbbr = splitline[2].strip()
            issn = splitline[3].strip()
            Publisher = splitline[7].strip()
            PublisherAddress = splitline[8].strip()
            Category = splitline[9].strip()

            set_ht = {}
            addtoset_ht = {}
            if FullTitle != "":
                set_ht["JCRTitle"] = FullTitle
                addtoset_ht["FullTitle"] = FullTitle
            if ISOAbbr != "":
                set_ht["JCRISOAbbr"] = ISOAbbr
                addtoset_ht["ISOAbbr"] = ISOAbbr
            if JCRAbbr != "":
                set_ht["JCRAbbr"] = JCRAbbr
            if issn != "" and issn != "****-****":
                addtoset_ht["ISSNSet"] = issn
            if Publisher != "":
                set_ht["JCRPublisher"] = Publisher
                addtoset_ht["Publisher"] = Publisher
            if PublisherAddress != "":
                set_ht["PublisherAddress"] = PublisherAddress
            if Category != "":
                Categories = Category.split("|")
                if len(Categories) == 1:
                    addtoset_ht["Category"] = Categories[0]
                else:
                    addtoset_ht["Category"] = {"$each": Categories}
            set_ht["JCR_Record"] = True

            # 使用ISSN核对
            result = MongoDB_Collection.count({"ISSNSet": issn})
            if result == 1:
                issn_count += 1
                MongoDB_Collection.update_one({"ISSNSet": issn}, {"$set": set_ht, "$addToSet": addtoset_ht},
                                              upsert=False)
                continue
            # 使用ISOAbbr核对
            if ISOAbbr != "":
                if isoabbr_ht.has_key(GetStringFeature(ISOAbbr)):
                    isoabbr_count += 1
                    MongoDB_Collection.update_one({"_id": isoabbr_ht[GetStringFeature(ISOAbbr)]},
                                                  {"$set": set_ht, "$addToSet": addtoset_ht},
                                                  upsert=False)
                    continue
            # 使用FullTitle核对
            if FullTitle != "":
                if title_ht.has_key(GetStringFeature(FullTitle)):
                    title_count += 1
                    MongoDB_Collection.update_one({"_id": title_ht[GetStringFeature(FullTitle)]},
                                                  {"$set": set_ht, "$addToSet": addtoset_ht},
                                                  upsert=False)
                    continue
            # 对不上的，直接插入数据
            upsert_count += 1
            MongoDB_Collection.update_one({},
                                          {"$set": set_ht, "$addToSet": addtoset_ht},
                                          upsert=True)
    infile.close()
    # print len(title_ht)
    # print len(isoabbr_ht)

    MongoDB_Collection.create_index([("FullTitle", ASCENDING)])
    MongoDB_Collection.create_index([("PubMedTitle", ASCENDING)])
    MongoDB_Collection.create_index([("PMCTitle", ASCENDING)])
    MongoDB_Collection.create_index([("JCRTitle", ASCENDING)])
    MongoDB_Collection.create_index([("PubMedISOAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("JCRISOAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("ISOAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("PubMedMedAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("PMCMedAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("MedAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("JCRAbbr", ASCENDING)])
    MongoDB_Collection.create_index([("PMCPublisher", ASCENDING)])
    MongoDB_Collection.create_index([("JCRPublisher", ASCENDING)])
    MongoDB_Collection.create_index([("Publisher", ASCENDING)])
    MongoDB_Collection.create_index([("PubMed_Record", ASCENDING)])
    MongoDB_Collection.create_index([("PMC_Record", ASCENDING)])
    MongoDB_Collection.create_index([("JCR_Record", ASCENDING)])

    logging.info("Update " + str(issn_count) + " records with ISSN.")
    logging.info("Update " + str(isoabbr_count) + " records with ISOAbbr.")
    logging.info("Update " + str(title_count) + " records with FullTitle.")
    logging.info("Insert " + str(upsert_count) + " new records.")
    logging.info("Finish analyzing JCR file.")


def GenerateUniquePublisher(MongoDB_Collection):
    logging.info("Start generating UniquePublisher.")
    results = MongoDB_Collection.find({"Publisher": {"$exists": 1}})
    publisher = ""
    for result in results:
        id = result["_id"]
        if result.has_key("JCRPublisher"):
            publisher = result["JCRPublisher"]
        else:
            publisher = result["PMCPublisher"]
        MongoDB_Collection.update_one({"_id": id}, {"$set": {"UniquePublisher": publisher}}, upsert=False)
    MongoDB_Collection.create_index([("UniquePublisher", ASCENDING)])
    logging.info("Finish generating UniquePublisher.")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    MongoDB_Connection, MongoDB_Collection = InitMongoDB()
    pubmed_filepath = "D:/data/journal/J_Entrez_20160803"
    InsertPubmed(MongoDB_Collection, pubmed_filepath)
    UpsertPMC(MongoDB_Collection, "D:/data/journal/jlist_20160803.csv")
    UpsertJCR(MongoDB_Collection, "D:/data/journal/jcr_2015.txt")
    GenerateUniquePublisher(MongoDB_Collection)
    CloseMongoDB(MongoDB_Connection)
